/*
This is the reference code of SolverChallenge to optimize by the participants
@Code version: 1.0
@Update date: 2021/5/17
@Author: Dechuang Yang, Haocheng Lian

Added some parallel packages: jxpamg
@Update date: 2021/5/28
@Author: Li Zhao
*/

typedef long double Real;
typedef int Int;
// typedef long long int Int;

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <sys/time.h>
#include <time.h> 

#include "mmio_highlevel.h"
#include "my_solver.h"
#include "mpi.h"


Real check_correctness(Int n, Int *row_ptr, Int *col_idx, Real *val, Real *x, Real *y, Real *b, Real *Res);
void store_x(Int n, Real *x, char *filename);
void load_b(Int n, Real *b, char *filename);

Int main(Int argc, char **argv)
{
    Int     m, n, nnzA, isSymmetricA;
    Int     *row_ptr; // the csr row pointer array of matrix A
    Int     *col_idx; // the csr column index array of matrix A
    Real  *val;  // the csr value array of matrix A
    Real  *x, *b; 
    Real  *y; 
    Real    norm_b;

    char    *filename_matrix = argv[1];  // the filename of matrix A
    char    *filename_b = argv[2];       // the filename of right-hand side vector b
    char    *filename_x = argv[3];       // the filename of solution vector x
    char    *matrix_tolerance = argv[4]; // the tolerance of input matrix
    // char    *sid = argv[5];              // solver id
    Int     iter = 0;                    // the number of iterations
    Real  residule = 0.0, RRN;
    Real  tolerance = atof(matrix_tolerance);
    Real  coff = 1;

    int     rank;
    int     size;
    Int     i;
    // Int     solver_id   = atoi(sid);
    Int     solver_id   = 12;
    Int     precond_id  = 3;
    Int     num_functions = 1;
    Int     max_iter = 500;
    
    // system("free -h");
    MPI_Init(NULL,NULL); 
    MPI_Comm_rank(MPI_COMM_WORLD, &rank); 
    MPI_Comm_size(MPI_COMM_WORLD, &size); 
    
    //-----------------------
    //  命令行修改参数
    //-----------------------
    Int arg_index = 1;
    while (arg_index < argc)
    {
        if ( strcmp(argv[arg_index], "-sid") == 0 )
        {
            arg_index ++;
            solver_id = atoi(argv[arg_index++]);
        }
        else if ( strcmp(argv[arg_index], "-maxiter") == 0 )
        {
            arg_index ++;
            max_iter = atoi(argv[arg_index++]);
        }
        else if ( strcmp(argv[arg_index], "-nf") == 0 )
        {
            arg_index ++;
            num_functions = atoi(argv[arg_index++]);
        }else{
            arg_index ++;
        }
    }
    if (rank == 0)
    {
        printf("===================================================================\n");
        printf("A  : %s\n", filename_matrix);
        printf("b  : %s\n", filename_b);
        printf("x  : %s\n", filename_x);
        printf("tol: %s\n", matrix_tolerance);
        printf("sid: %d\n", solver_id);
        printf("nf : %d\n", num_functions);
        printf("===================================================================\n");

        printf("\n\n+++++++++++++++++++++ MPI using %d processors +++++++++++++++++++++\n\n", size);

        //load matrix
        printf("load matrix:\n");
        mmio_allinone(&m, &n, &nnzA, &isSymmetricA, &row_ptr, &col_idx, &val, filename_matrix);
        printf("matrix info: m: %lld, n: %lld, nnz: %lld, isSymmertirc: %lld\n", m, n, nnzA, isSymmetricA);

        if (m != n)
        {
            printf("Invalid matrix size.\n");
            return 0;
        }

        x = (Real *)malloc(sizeof(Real) * n);
        y = (Real *)malloc(sizeof(Real) * m);
        b = (Real *)malloc(sizeof(Real) * m);

        // load right-hand side vector b
        //printf("load right-hand side vector:\n");
        load_b(n, b, filename_b);

        norm_b = vec2norm(b, n);
        //printf("the norm of b is %Lf\n", norm_b);
        // load inital guess vector x
        //printf("load inital guess vector x:\n");
        load_b(n, x, filename_x);
        
        //for (i = 0; i < n; i++) x[i] = 0; // 初始化为 0 

        spmv(n, row_ptr, col_idx, val, x, y);  // y=A*x_0  
        Int i;
        Real *r0 = (Real *)malloc(sizeof(Real) * n);
        for (i = 0; i < m; i++)
        {
            r0[i] = y[i] - b[i];
        }
        Real r0_L2 = vec2norm(r0, n);    // r0_L2
        //printf("the norm of r0 is %Lf\n", r0_L2);
    }



    // solve x and record wall-time
    struct timeval t_start, t_stop;
    gettimeofday(&t_start, NULL);

    //-------------------------------------------------------------------------
    //  boomer_amg
    //-------------------------------------------------------------------------  
    if (rank == 0) printf("\nJXPAMG_Solver_Interface(int, long double -- version):\n");    
    JXPAMG_Solver_Interface(argc, argv, n, row_ptr, col_idx, val, x, b, tolerance, solver_id, num_functions, &iter); //确认收敛是否是除以的r_0
    

    gettimeofday(&t_stop, NULL);
    Real total_time = (t_stop.tv_sec - t_start.tv_sec) * 1000.0 + (t_stop.tv_usec - t_start.tv_usec) / 1000.0;
    
    if (rank == 0)
    {
        //store x to a file
        // store_x(n, x, filename_x);
        store_x(n, x, "x.txt");

        //check the correctness
        //print the #iteration, residual and total_time
        RRN = check_correctness(n, row_ptr, col_idx, val, x, y, b, &residule);
        printf("\n\n===============================================================================\n");
        printf("iteration = %d, residual = %Le, RRN = %Le, total_time = %.5Lf sec\n", iter, residule, RRN, total_time / 1000.0);
        printf("===============================================================================\n");
    }
    
    ///
    MPI_Finalize();
    return 0;
}

//validate the x by b-A*x
Real check_correctness(Int n, Int *row_ptr, Int *col_idx, Real *val, Real *x, Real *y, Real *b, Real *Res)
{
    Real *b_new = (Real *)malloc(sizeof(Real) * n);
    Real *check_b = (Real *)malloc(sizeof(Real) * n);
    Real *r0 = (Real *)malloc(sizeof(Real) * n);
    spmv(n, row_ptr, col_idx, val, x, b_new);  // b_new=A*x_out
    Int i;
    for (i = 0; i < n; i++)
    {
        check_b[i] = b_new[i] - b[i];       // check_b = r
        r0[i] = y[i] - b[i];
    }   
    *Res =  vec2norm(check_b, n);   // Res = rnorm
    Real r0_L2 = vec2norm(r0, n);    // r0_L2
    Real b_L2 = vec2norm(b, n);     // b_L2 = bnorm
    // printf("\nZL: ||b-Ax||_2: %Le, ||b||_2: %Le\n", vec2norm(check_b, n), b_L2);
    
    //return *Res / b_L2;
    return *Res / r0_L2;
}

//store x to a file
void store_x(Int n, Real *x, char *filename)
{
    FILE *p = fopen(filename, "w");
    fprintf(p, "%lld\n", n);
    Int i;
    for (i = 0; i < n; i++)
        fprintf(p, "%Lf\n", x[i]);
    fclose(p);
}

//load right-hand side vector b
void load_b(Int n, Real *b, char *filename)
{
    FILE *p = fopen(filename, "r");
    Int n_right;
    Int r = fscanf(p, "%lld", &n_right);
    // printf("n_right: %lld, n: %lld.\n", n_right, n);
    if (n_right != n)
    {
        fclose(p);
        printf("Invalid size of b.\n");
        return;
    }
    Int i;
    for (i = 0; i < n_right; i++)
        r = fscanf(p, "%Lf", &b[i]);
    fclose(p);
}
